import argparse
import os
import torch

def merge_checkpoints(encoder_dir, decoder_dir, final_dir, tp, pp):
    # Merging parameters from encoder and decoder directories to final directory
    for tp_idx in range(tp):
        for pp_idx in range(pp):
            if pp == 1:
                # For pp == 1, folder format is different
                decoder_file = f"{decoder_dir}/iter_0000001/mp_rank_{tp_idx:02d}/model_optim_rng.pt"
                final_output_dir = f"{final_dir}/iter_0000001/mp_rank_{tp_idx:02d}"
                # Encoder file only needs to be merged into PP rank 0
                encoder_file = f"{encoder_dir}/iter_0000001/mp_rank_{tp_idx:02d}/model_optim_rng.pt"
                encoder_ckp = torch.load(encoder_file, map_location='cpu')
            else:
                if pp_idx == 0:
                    # Encoder file only needs to be merged into PP rank 0
                    encoder_file = f"{encoder_dir}/iter_0000001/mp_rank_{tp_idx:02d}/model_optim_rng.pt"
                    encoder_ckp = torch.load(encoder_file, map_location='cpu')
                else:
                    encoder_ckp = None
                
                decoder_file = f"{decoder_dir}/iter_0000001/mp_rank_{tp_idx:02d}_{pp_idx:03d}/model_optim_rng.pt"
                final_output_dir = f"{final_dir}/iter_0000001/mp_rank_{tp_idx:02d}_{pp_idx:03d}"
            
            final_output_file = f"{final_output_dir}/model_optim_rng.pt"

            # Create final output directory if it does not exist
            os.makedirs(final_output_dir, exist_ok=True)

            # Load and merge parameters from decoder checkpoints (and encoder if PP rank is 0)
            decoder_ckp = torch.load(decoder_file, map_location='cpu')
            if encoder_ckp is not None:
                merged_ckp = decoder_ckp
                merged_ckp['model'].update(encoder_ckp['model'])
            else:
                merged_ckp = decoder_ckp
            
            
            torch.save(merged_ckp, final_output_file)

            print(f"Merged checkpoint saved to {final_output_file}")

    # Update latest_checkpointed_iteration.txt
    latest_iteration_file = os.path.join(final_dir, "latest_checkpointed_iteration.txt")
    with open(latest_iteration_file, "w") as f:
        f.write("1")

    print("All checkpoints have been merged successfully.")

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Merge encoder and decoder Megatron checkpoints.")
    parser.add_argument('--encoder-dir', type=str, required=True, help="Path to the Megatron encoder directory.")
    parser.add_argument('--decoder-dir', type=str, required=True, help="Path to the Megatron decoder directory.")
    parser.add_argument('--final-dir', type=str, required=True, help="Path to save the merged Megatron checkpoint.")
    parser.add_argument('--tp', type=int, required=True, help="Tensor parallel size.")
    parser.add_argument('--pp', type=int, required=True, help="Pipeline parallel size.")

    args = parser.parse_args()

    merge_checkpoints(args.encoder_dir, args.decoder_dir, args.final_dir, args.tp, args.pp)
